import random
import pandas as pd
from cadCAD.configuration import Experiment
from cadCAD.configuration.utils import config_sim
from cadCAD.engine import ExecutionMode, ExecutionContext, Executor
from cadCAD import configs

def state_update_func1(_params, substep, sH, s, _input):
    var1 = _params['var1']
    new_state = var1 * 2
    return 'state_var1', new_state


def state_update_func2(_params, substep, sH, s, _input):
    var2 = _params['var2']
    new_state = var2 ** 2
    return 'state_var2', new_state


def get_psubs():
    PSUBs = [
        {
            "policies": {},
            'variables': {
                'state_var1': state_update_func1,
                'state_var2': state_update_func2,
            }
        },
    ]
    return PSUBs


# 参数
def getVar():
    var_lists = {
        'var1': [1, 2, 3],
        'var2': [5, 10, 15]
    }
    return var_lists


if __name__ == "__main__":
    var_lists = getVar()
    print('variable lists', var_lists)

    exp = Experiment()
    exp.append_model(
        model_id='test1',
        initial_state={
            'state_var1': 1,
            'state_var2': 2,
        },
        sim_configs=config_sim({
            'N': 1,
            'T': range(1),
            'M': var_lists
        }),
        partial_state_update_blocks=get_psubs()
    )

    exec_mode = ExecutionMode()
    exec_context = ExecutionContext(context=exec_mode.multi_proc)
    simulation = Executor(exec_context=exec_context, configs=exp.configs)
    raw_system_events, tensor_field, sessions = simulation.execute()

    simulation_result = pd.DataFrame(raw_system_events)
    print(simulation_result)

    # print()
    # print("Tensor Field: ")
    # print(tensor_field)
    # # print(tabulate(tensor_field, headers='keys', tablefmt='psql'))
    # print("Result: System Events DataFrame")
    # # print(tabulate(simulation_result, headers='keys', tablefmt='psql'))
    # print(simulation_result)
    # print()
